"""
直接启动Isaac Sim并加载G1机器人的完整脚本 - 重构版本
集成键盘控制和ONNX策略模型，实现完整的机器人控制
使用模块化架构，便于扩展和维护
"""

import os
import sys
import numpy as np
import carb
import omni
import torch

# 1. 必须首先初始化SimulationApp
from isaacsim import SimulationApp
simulation_app = SimulationApp({
    'headless': False,
    'anti_aliasing': 0,
    'hide_ui': False,
    'multi_gpu': False,
    'use_fabric': False,
    'physics_dt': '1/240',
    'rendering_dt': '1/240',
})

# 2. 设置物理参数（必须在导入其他Isaac Sim模块之前）
simulation_app._carb_settings.set('/physics/cooking/ujitsoCollisionCooking', False)
print("设置物理参数：ujitsoCollisionCooking = False（与源代码一致）")
print("设置时间步长：physics_dt=1/240, rendering_dt=1/240（与源代码一致）")
print("设置use_fabric=False（与源代码一致）")

# 3. 启用必要的扩展
from omni.isaac.core.utils.extensions import enable_extension
enable_extension("omni.isaac.dynamic_control")
print("✅ 已启用Dynamic Control扩展: omni.isaac.dynamic_control")

# 4. 现在可以安全导入Isaac Sim模块
from omni.isaac.core import World
from omni.isaac.core.utils.stage import add_reference_to_stage
from omni.isaac.core.articulations import Articulation

# 4. 导入项目模块
from graduation_project.core import ASSET_PATH
from graduation_project.core.config import ROBOT_CONFIG, G1_PD_GAINS, SIMULATION_LOOP
from graduation_project.core.robot import SimpleG1Robot
from graduation_project.controllers.keyboard import DirectKeyboardController
from graduation_project.controllers.g1_controller import G1MoveBySpeedController
from graduation_project.utils.joint_mapping import JointMapping


def main():
    """主函数 - 线性执行流程"""
    # 5. 创建物理世界
    print("创建物理世界...")
    world = World(
        physics_dt=1/240,
        rendering_dt=1/240,
        stage_units_in_meters=1.0,
        sim_params={'use_fabric': False}
    )

    # 6. 加载空场景
    print("正在加载空场景...")
    add_reference_to_stage(prim_path="/World/scene", usd_path=ROBOT_CONFIG['empty_scene_path'])

    # 7. 加载G1机器人
    print("加载G1机器人...")
    add_reference_to_stage(prim_path=ROBOT_CONFIG['prim_path'], usd_path=ROBOT_CONFIG['g1_usd_path'])
    raw_robot = Articulation(prim_path=ROBOT_CONFIG['prim_path'], name="g1")
    # raw_robot.set_enabled_self_collisions(False)  # 禁用自碰撞

    # 8. 设置机器人初始位置并初始化
    print("设置机器人初始位置...")
    raw_robot.set_world_pose(position=np.array(ROBOT_CONFIG['initial_position']))
    world.scene.add(raw_robot)
    world.reset()

    # 9. 创建机器人包装类
    print("创建机器人包装类...")
    g1_robot = SimpleG1Robot(raw_robot)

    # 10. 等待物理稳定
    # print("等待物理稳定...")
    # for i in range(10):
    #     world.step(render=True)

    # 11. 创建刚体映射
    print("创建刚体映射...")
    g1_robot.create_rigid_bodies()
    g1_robot.post_reset()
    print(f"刚体映射创建完成，映射包含 {len(g1_robot._rigid_body_map)} 个刚体")

    # 12. 设置关节PD参数
    num_dofs = g1_robot.num_dof
    print(f"G1机器人自由度数量: {num_dofs}")

    joint_mapping = JointMapping()
    kps = joint_mapping.gym2sim(np.array(G1_PD_GAINS['kps_gym']))
    kds = joint_mapping.gym2sim(np.array(G1_PD_GAINS['kds_gym']))
    g1_robot.set_gains(kps=kps[:num_dofs], kds=kds[:num_dofs])
    print(f"关节PD参数设置完成")

    # 13. 设置物理引擎参数
    # g1_robot.articulation.set_solver_position_iteration_count(4)
    # g1_robot.articulation.set_solver_velocity_iteration_count(0)
    # print("已设置物理引擎关键参数：position_iterations=4, velocity_iterations=0")

    # 14. 设置默认关节位置
    # print("设置默认站立姿势...")
    # default_dof_pos_gym = np.concatenate([
    #     np.array([-0.1, 0.0, 0.0, 0.3, -0.2, 0.0, -0.1, 0.0, 0.0, 0.3, -0.2, 0.0]),  # lower joints
    #     np.zeros(15),  # upper joints
    # ])
    # default_dof_pos_sim = joint_mapping.gym2sim(default_dof_pos_gym)

    # if num_dofs >= 27:
    #     all_positions = np.zeros(num_dofs)
    #     all_positions[:12] = default_dof_pos_sim[:12]
    #     g1_robot.set_joint_positions(all_positions)
    #     print(f"设置所有27个关节：下肢12个+上肢15个零位置")

    # 15. 初始化控制器
    print("正在初始化控制模块...")
    keyboard_controller = DirectKeyboardController()
    g1_controller = G1MoveBySpeedController(model_path=ROBOT_CONFIG['onnx_model_path'], robot=g1_robot)
    g1_robot.controllers['move_by_speed'] = g1_controller
    print("控制器已添加到机器人系统")

    # 16. 打印使用说明
    print("\n=== G1机器人键盘控制 ===")
    print("使用以下按键控制机器人移动：")
    print("  I键: 前进")
    print("  K键: 后退")
    print("  J键: 左移")
    print("  L键: 右移")
    print("  U键: 左转")
    print("  O键: 右转")
    print("========================\n")

    # 17. 主控制循环
    print("开始主控制循环...")
    frame_count = 0
    debug_counter = 0
    render_interval = SIMULATION_LOOP['render_interval']
    render_trigger = 0

    while simulation_app.is_running() and frame_count < SIMULATION_LOOP['max_frames']:
        # 获取键盘输入
        command = keyboard_controller.command
        x_speed = command[0] - command[1]  # I - K
        y_speed = command[2] - command[3]  # J - L
        z_speed = command[4] - command[5]  # U - O

        debug_counter += 1

        # 调试信息
        if debug_counter % SIMULATION_LOOP['debug_interval'] == 0:
            print(f"\n--- 控制器调试信息 (帧数: {frame_count}) ---")
            print(f"键盘命令: [{x_speed:.3f}, {y_speed:.3f}, {z_speed:.3f}]")

            current_joints = g1_robot.get_joint_positions()
            joint_velocities = g1_robot.get_joint_velocities()
            print(f"当前关节位置范围: [{np.min(current_joints):.3f}, {np.max(current_joints):.3f}]")
            print(f"关节速度范围: [{np.min(joint_velocities):.3f}, {np.max(joint_velocities):.3f}]")

            controller = g1_robot.controllers['move_by_speed']
            print(f"控制器 apply_times_left: {controller._apply_times_left}")

        # 速度限制
        max_speed = 1.0
        x_speed = np.clip(x_speed, -max_speed, max_speed)
        y_speed = np.clip(y_speed, -max_speed, max_speed)
        z_speed = np.clip(z_speed, -max_speed, max_speed)

        # 应用动作
        action = {
            'move_by_speed': [float(x_speed), float(y_speed), float(z_speed)]
        }
        g1_robot.apply_action(action)

        # 渲染控制
        render_trigger += 1
        need_render = render_trigger > render_interval
        if render_trigger > render_interval:
            render_trigger = 0

        # 执行仿真步骤
        world.step(render=need_render)
        frame_count += 1

        # 状态信息
        if frame_count % SIMULATION_LOOP['status_interval'] == 0:
            obs = g1_robot.get_obs()
            position = obs['position']
            orientation = obs['orientation']

            # 计算欧拉角
            x, y, z, w = orientation
            yaw = np.arctan2(2.0 * (w * z + x * y), 1.0 - 2.0 * (y * y + z * z))
            pitch = np.arcsin(2.0 * (w * y - x * z))
            roll = np.arctan2(2.0 * (w * x + y * z), 1.0 - 2.0 * (x * x + y * y))

            print(f"帧数: {frame_count:4d} [G1Controller] | "
                    f"位置: [{position[0]:7.3f}, {position[1]:7.3f}, {position[2]:7.3f}] | "
                    f"姿态: [{np.degrees(roll):6.1f}°, {np.degrees(pitch):6.1f}°, {np.degrees(yaw):6.1f}°]")

    print("G1Robot控制器测试完成")


    print("关闭Isaac Sim...")
    simulation_app.close()
    print("仿真已结束")


if __name__ == "__main__":
    main()